package goalie;

import goalie.util.Pair;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStreamWriter;
import java.net.URL;
import java.util.ArrayList;
import java.util.Date;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.Map.Entry;

// TODO: allow for stages that don't run, but only set values (such as 'global')
// TODO: look again into outputs that are at the same time inputs/lets, i.e.
//  make sure they ahve to be the same always
public class Workflow implements Component {
  public static String REFERRED_NAME = "workflow";
  // names of instructions
  private static String INSTR_RUN = "run";
  private static String INSTR_LET = "let";
  private static String INSTR_INPUT = "input";
  // input/output names
  public static String CONFIG = "config";
  public static String RESULT = "result";
  // All components for Workflow.
  private static Map<String, Class<? extends Component>> name2component = 
    new HashMap<String, Class<? extends Component>>();

  Map<String, Set<String>> task2dependsOn = new HashMap<String, Set<String>>();
  // This contains the information from the config file, keyed by task
  // task -> {inputname, value}
  Map<String, Map<String, String>> task2inputs = 
    new HashMap<String, Map<String, String>>();
  Map<String, Map<String, String>> task2lets = 
    new HashMap<String, Map<String, String>>();

  // TODO: This stores output values specified in config file. Tasks with
  // outputs in config file are not run again -- they are assume to be ready.
  // Map<String, Map<String, String>> task2outputs;

  Map<String, Component> task2component = new HashMap<String, Component>();
  Set<String> goals = new HashSet<String>();

  public static void registerCommon() {
    register("goalie.Workflow");
    registerAll("goalie.common");
  }
  

  /**
   * Deletes the resulting file ("result").
   */
  @Override
  public void cleanup(Map<String, String> outputs) {
    String fileName = outputs.get(RESULT);
    File f = new File(fileName);
    f.delete();
  }

  /**
   * This runs a workflow with a given properties file. Returns a result file.
   * 
   */
  public String run(String inputFN) throws IOException {
    Map<String, String> params = new HashMap<String, String>();
    params.put(CONFIG, inputFN);
    return run(params).get(RESULT);
  }

  /**
   * This runs a workflow.
   * Input:
   * "config" : The properties file with the Workflow configuration.
   * Output:
   * "result" : A properties file with the results of the goal components.
   */
  @Override
  public Map<String, String> run(Map<String, String> params) throws 
      IOException {
    String configFN = params.get(CONFIG);
    Properties props = new Properties();
    InputStream is = new FileInputStream(configFN);
    props.load(is);
    parseProperties(props);
    List<String> runOrder = taskRunOrder();

    System.out.println("Task dependencies:");
    // To keep track when cleanup of required tasks has to be called.
    Map<String, Integer> task2numDependents = new HashMap<String, Integer>();
    for (String dependent : runOrder) {
      System.out.print(dependent);
      for (String required : task2dependsOn.get(dependent)) {
        int numDependents = task2numDependents.containsKey(required) ? 
            task2numDependents.get(required) : 0;
        task2numDependents.put(required, numDependents + 1);
        System.out.print(" " + required);
      }
      System.out.println();
    }
    System.out.println();

    Map<String, Map<String, String>> task2outputs = 
      new HashMap<String, Map<String, String>>();

    for (String task : runOrder) {
      Component comp = task2component.get(task);
      System.out.print(task + " RUNNING");
      Date start = new Date();
      Map<String, String> outputs = comp.run(params(task, task2outputs));
      Date end = new Date();
      System.out.print(" DONE");
      System.out.println(" (" + timeElapsedString(start, end) + ")");
      task2outputs.put(task, outputs);
      for (String required : task2dependsOn.get(task)) {
        task2numDependents.put(required, task2numDependents.get(required) - 1);
        if (task2numDependents.get(required) == 0 && 
            !goals.contains(required)) {
          Map<String, String> toClean = task2outputs.get(required);
          System.out.print(required + " CLEANUP");
          task2component.get(required).cleanup(toClean);
          System.out.println(" DONE");
        }
      }
    }

    String resultFN = configFN + ".result";
    Map<String, String> resultMap = new HashMap<String, String>();
    resultMap.put(RESULT, resultFN);
    writeGoalOutputs(task2outputs, resultFN);
    System.out.println("\nResult: " + resultFN);
    return resultMap;
  }

  private String timeElapsedString(Date begin, Date end) {
    long diff = end.getTime() - begin.getTime();
    long secondInMillis = 1000;
    long minuteInMillis = secondInMillis * 60;
    long hourInMillis = minuteInMillis * 60;
    long dayInMillis = hourInMillis * 24;

    StringBuffer sb = new StringBuffer();
    long days = diff / dayInMillis;
    if (days > 0) sb.append(days + "d ");
    diff = diff % dayInMillis;
    long hours = diff / hourInMillis;
    if (hours > 0) sb.append(hours + "h ");
    diff = diff % hourInMillis;
    long minutes = diff / minuteInMillis;
    if (minutes > 0) sb.append(minutes + "m ");
    diff = diff % minuteInMillis;
    double secs = diff / (float) secondInMillis;
    sb.append(String.format(Locale.US,"%1$,.2fs", secs));
    return sb.toString();
  }


  /**
   * Writes the outputs for the goals to a file.
   * 
   * @param task2outputs
   * @param resultFN
   * @throws IOException
   */
  private void writeGoalOutputs(Map<String, Map<String, String>> task2outputs,
      String resultFN) throws IOException {
    BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(
        new FileOutputStream(resultFN), "UTF-8"));
    for (String task : task2outputs.keySet()) {
      if (goals.contains(task)) {
        for (Entry<String, String> argVal : 
            task2outputs.get(task).entrySet()) {
          bw.append("# " + task);
          bw.newLine();
          bw.append(task + "." + argVal.getKey());
          bw.append(" = output ");
          bw.append(argVal.getValue());
        }
        bw.newLine();
      }
    }
    bw.flush();
    bw.close();
  }

  /**
   * This processes argument definitions for a task and returns the argument
   * map, i.e. a fully expanded form of the argument value.
   * 
   * Eg. if there are two task with configurations:
   * 
   * task0.a = let arg1 
   * task1.c = let arg3 
   * task1.d = input task0.a task0.b task1.c 
   * task1.e = input task1.c task0.b task0.a
   * 
   * then params("task1", {"task0"->{"b"->"arg2"}}) returns:
   * 
   * {"d"->"arg1 arg2 arg3", "e"->"arg3 arg2 arg1"}
   * 
   * @param task
   *          the task for which to expand arguments
   * @param task2outputs
   *          the outputs of previously run tasks
   * @return a map with expanded arguments
   */
  private Map<String, String> params(String task,
      Map<String, Map<String, String>> task2outputs) {
    // Maps a particular input parameter for task to a string with the value.
    Map<String, String> paramMap = new HashMap<String, String>();
    // Values specified as 'let' are just taken over.
    for (Entry<String, String> let : task2lets.get(task).entrySet()) {
      paramMap.put(let.getKey(), let.getValue());
    }
    // Values specified as 'input' are replaced by the output of the respective
    // tasks, or by the respective 'let' value.
    for (Entry<String, String> input : task2inputs.get(task).entrySet()) {
      String inputStr = "";
      // An input value can consist of several specifications, eg.:
      // task2.valA = input task1.set_valA task1.ouput_valB task2.some_set_val
      for (Pair<String, String> tv : taskVals(input.getValue())) {
        boolean found = false;
        // Is it an output from another task?
        if (task2outputs.containsKey(tv.first)
            && task2outputs.get(tv.first).containsKey(tv.second)) {
          inputStr += task2outputs.get(tv.first).get(tv.second);
          found = true;
        } else {
          for (Entry<String, String> letval : task2lets.get(tv.first).entrySet()) {
            // Is it specified as 'let'?
            if (letval.getKey().equals(tv.second)) {
              inputStr += letval.getValue();
              found = true;
              break;
            }
          }
        }
        if (found) {
          inputStr += " ";
        } else {
          throw new IllegalStateException(tv.first + "." + tv.second
              + " not specified.");
        }
      }
      paramMap.put(input.getKey(), inputStr.trim());
    }
    return paramMap;
  }

  /**
   * Returns a running order so that the goals and all tasks they depend on are 
   * run one after each other.
   * 
   * @return
   */
  private List<String> taskRunOrder() {
    Set<String> marked = new HashSet<String>();
    List<String> order = new ArrayList<String>();
    for (String g : goals) {
      visit(g, marked, order);
    }
    return order;
  }

  /**
   * This visits a task node in the dependency graph and adds it to the running
   * order. Changes are reflected in the arguments marked and order (i.e. this
   * method is destructive in the arguments).
   */
  private void visit(String task, Set<String> marked, List<String> order) {
    if (!marked.contains(task)) {
      marked.add(task);
      for (String prereq : task2dependsOn.get(task)) {
        visit(prereq, marked, order);
      }
      order.add(task);
    }
  }

  /**
   * This processes the given Properties and adds the information to the 
   * Workflow.
   * 
   * @param props
   */
  private void parseProperties(Properties props) {
    for (String p : props.stringPropertyNames()) {
      if (p.equals("goal")) {
        for (String g : props.getProperty(p).trim().split("\\s")) {
          goals.add(g);
        }
        continue;
      }
      String[] taskAndInput = p.trim().split("\\.", 2);
      String[] typeAndVal = props.getProperty(p).trim().split("\\s", 2);

      if (typeAndVal.length != 2
          || !(taskAndInput.length == 1 || taskAndInput.length == 2)) {
        throw new IllegalArgumentException("Ill-formed property: " + p);
      }
      String type = typeAndVal[0];
      String val = typeAndVal[1];
      String task = taskAndInput[0];

      if (!task2dependsOn.containsKey(task)) {
        task2dependsOn.put(task, new HashSet<String>());
        assert (!task2inputs.containsKey(task));
        task2inputs.put(task, new HashMap<String, String>());
        assert (!task2lets.containsKey(task));
        task2lets.put(task, new HashMap<String, String>());
      }
      if (type.equals(INSTR_RUN)) {
        if (taskAndInput.length != 1) {
          throw new IllegalArgumentException("Ill-formed property: " + p);
        }
        try {
          task2component.put(task, name2component.get(val).newInstance());
        } catch (Exception e) {
          e.printStackTrace();
          throw new IllegalStateException("Could not instatiate task " + val);
        }
      } else if (type.equals(INSTR_LET)) {
        if (taskAndInput.length != 2) {
          throw new IllegalArgumentException("Ill-formed property: " + p);
        }
        task2lets.get(task).put(taskAndInput[1], val);
      } else if (type.equals(INSTR_INPUT)) {
        if (taskAndInput.length != 2) {
          throw new IllegalArgumentException("Ill-formed property: " + p);
        }
        task2inputs.get(task).put(taskAndInput[1], val);
      } else {
        throw new IllegalArgumentException("Ill-formed property: " + p);
      }
    }
    
    for (String taskWithInput : task2inputs.keySet()) {
      // For checking and deopendency construction only consider tasks that
      // actually have input.
      if (task2inputs.get(taskWithInput).isEmpty()) continue;
      if (!task2component.containsKey(taskWithInput)) {
        throw new IllegalArgumentException("Task " + taskWithInput + 
            " has input but is not runnable.");
      }
      for (Entry<String, String> inputAndVal : 
          task2inputs.get(taskWithInput).entrySet()) {
        for (Pair<String, String> prereqAndOutput : 
            taskVals(inputAndVal.getValue())) {
          String preTask = prereqAndOutput.first;
          if ( task2lets.containsKey(preTask) && 
              task2lets.get(preTask).containsKey(prereqAndOutput.second) ) {
            // Nothing to do.
          } else if (!taskWithInput.equals(preTask) &&
                      task2component.containsKey(preTask)) {
            task2dependsOn.get(taskWithInput).add(preTask);            
          } else {
            throw new IllegalArgumentException(
              "Task input must be output from runnable or set as 'let'.");
          }
        }
      }
    }
  }

  /**
   * This returns, for a string of input specifications (the part after "input"
   * in "taskX.argX = input task0.arg0 task1.arg1") the pairs of task and
   * argument names.
   * 
   * @param cfgStr
   * @return
   */
  private static List<Pair<String, String>> taskVals(String cfgStr) {
    List<Pair<String, String>> l = new ArrayList<Pair<String, String>>();
    for (String taskOp : cfgStr.trim().split("\\s")) {
      String[] taskAndOutput = taskOp.split("\\.", 2);
      if (taskAndOutput.length != 2) {
        throw new IllegalArgumentException("Ill-formed task input: " + cfgStr);
      }
      l.add(new Pair<String, String>(taskAndOutput[0], taskAndOutput[1]));
    }
    return l;
  }

  /**
   * This maps a short identifier to an instance of a component.
   * 
   * @param componentName
   * @param mustBeComponent Whether it is expected that the class with 
   *  componentName is really a Component (throws an Exception if not), or 
   *  whether a non-Component is gracefully ignored.
   */
  public static void register(String componentName, 
      boolean mustBeComponent) {
    try {
      Class<? extends Component> classToRegister = Class.forName(componentName)
          .asSubclass(Component.class);
      name2component.put(classToRegister.getField("REFERRED_NAME").get(null)
          .toString(), classToRegister);
    } catch (Exception e) {
      if (mustBeComponent) {
        e.printStackTrace();
        throw new IllegalStateException("Unable to register " + componentName);
      }
    }
  }

  /**
   * This registers a Component, so that it can be referred to by a Workflow 
   * configuration by its REFERRED_NAME.
   * 
   * @param componentName
   */
  public static void register(String componentName) {
    register(componentName, true);
  }

  /**
   * This registers all Components in the given package and sub-packages, so 
   * that they can be referred to by a Workflow configuration by their 
   * respective REFERRED_NAME's.
   * 
   * @param componentName
   */
  public static void registerAll(String packageName) {
    ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
    String path = packageName.replace('.', '/');
    try {
      Enumeration<URL> resources = classLoader.getResources(path);
      while (resources.hasMoreElements()) {
        URL resource = resources.nextElement();
        File directory = new File(resource.toURI());
        registerAll(directory, packageName);
      }
    } catch (Exception e) {
      e.printStackTrace();
      throw new IllegalArgumentException("Could not register package: "
          + packageName);
    }
  }

  /**
   * Recursively registers Components in a directory.
   * @param directory
   * @param packageName
   * @throws ClassNotFoundException
   * @throws InstantiationException
   * @throws IllegalAccessException
   */
  private static void registerAll(File directory, String packageName)
      throws ClassNotFoundException, InstantiationException,
      IllegalAccessException {
    for (File file : directory.listFiles()) {
      String fileName = file.getName();
      if (file.isDirectory()) {
        registerAll(file, packageName + "." + fileName);
      } else if (fileName.endsWith(".class") && !fileName.contains("$")) {
        register(packageName + '.'
            + fileName.substring(0, fileName.length() - 6), false);
      }
    }
  }
}
